import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.axisartist as axisartist


def square(t, T, delta):
    # 方波信号
    if t % T <= 2 * delta:
        return 1
    else:
        return 0

def triangle(t, T, delta):
    # 三角信号
    t += 2
    if t % T <= delta:
        return t % T / delta
    elif t % T <= 2 * delta:
        return 1 - (t - delta) % T / delta
    else:
        return 0
    
def set_axis(fig, number, title, x_range=(-10, 15), y_range=(-1, 1)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(x_range[0], x_range[1])
    ax.set_ylim(y_range[0], y_range[1])
    ax.set_title(title)
    
def correlate(tau, fun1, fun2, t, T, delta):
    return np.sum(np.array([fun1(i, T, delta) * fun2(i - tau, T, delta) for i in t]))


if __name__ == "__main__":
    
    fig = plt.figure(figsize=(12, 8))
    plt.rcParams['font.family'] = ['Arial Unicode MS']
    
    T = 8
    delta = 2
    t = np.arange(-4 * T, 4 * T, 1)

    fun1 = np.array([square(i, T, delta) for i in t])
    set_axis(fig, 221, "原始的方波信号", (-4 * T, 4 * T), (-1, 2))
    plt.stem(t, fun1)
    
    fun2 = np.array([triangle(i, T, delta) for i in t])
    set_axis(fig, 222, "原始的三角波信号", (-4 * T, 4 * T), (-1, 2))
    plt.stem(t, fun2)
    
    fun3 = np.array([correlate(i, square, triangle, t, T, delta) for i in t])
    set_axis(fig, 223, "$r_{xy}$", (-4 * T, 4 * T), (-1, 20))
    plt.stem(t, fun3)
    
    fun4 = np.array([correlate(i, triangle, square, t, T, delta) for i in t])
    set_axis(fig, 224, "$r_{yx}$", (-4 * T, 4 * T), (-1, 20))
    plt.stem(t, fun4)
    
    plt.show()